import ast
import os
import tempfile
import gradio as gr
from PIL import Image
import requests


async def upload_image(model, image):
    image = Image.fromarray(image)
    folder_name = f"{os.getcwd()}/temp"
    os.makedirs(folder_name, exist_ok=True)
    temp_file = os.path.basename(tempfile.NamedTemporaryFile(suffix=".png").name)
    save_path = os.path.join(folder_name, temp_file)
    image.save(save_path)
    if model == "DeeplabV3":
        with open(save_path, "rb") as img_file:
            files = {"image": (temp_file, img_file, "image/png")}
            response = requests.post("http://0.0.0.0:8000/upload_image/", files=files)
        os.remove(save_path)
        content = response.json()["time"]
        path = response.json()["save_path"]
        gr.Info("上传成功")
        state = gr.State(2)
        return content, path, state
    elif model == "nnunet":
        with open(save_path, "rb") as img_file:
            files = {"image": (temp_file, img_file, "image/png")}
            response = requests.post("http://114.55.245.149:7833/upload_image/", files=files)
        os.remove(save_path)
        content = response.json()["time"]
        path = f"http://114.55.245.149:7833/colorinput/{content}/{temp_file}"
        gr.Info("上传成功")
        return content, path
    else:
        return gr.Info("请选择模型")


async def process_img(model, timestamp, path):
    if model == "DeeplabV3":
        data = {"user_string": timestamp}
        requests.post("http://0.0.0.0:8000/process_image/", params=data)
        base_image_name = os.path.basename(path)
        data1 = {"user_string": timestamp, "base_image_name": base_image_name}
        requests.post("http://0.0.0.0:8000/mix_together/", params=data1)
        cal_out = f"http://0.0.0.0:8000/output/{timestamp}/cal/cal.png"
        blended_img = f"http://0.0.0.0:8000/output/{timestamp}/cal/blended_image.png"
        return cal_out, blended_img
    elif model == "nnunet":
        data = {"user_string": timestamp}
        requests.post("http://114.55.245.149:7833/process_image/", params=data)
        base_image_name = os.path.basename(path)
        data1 = {"user_string": timestamp, "base_image_name": base_image_name}
        requests.post("http://114.55.245.149:7833/mix_together/", params=data1)
        cal_out = f"http://114.55.245.149:7833/colorinput/{timestamp}/{timestamp}/color/color.png"
        blended_img = f"http://114.55.245.149:7833/colorinput/{timestamp}/blended_image.png"
        return cal_out, blended_img


async def mix_img(model, timestamp, path):
    base_image_name = os.path.basename(path)
    if model == "DeeplabV3":
        paths = []
        data = {"user_string": timestamp}
        data1 = {"user_string": timestamp, "base_image_name": base_image_name}
        response = requests.post("http://0.0.0.0:8000/mix_images/", params=data1)
        #response = requests.get("http://114.55.245.149:6655/list_images/", params=data)
        img_list = response.text
        image_list = ast.literal_eval(img_list)
        for image_name in image_list:
            file_path = os.path.join(f'http://0.0.0.0:8000/output/{timestamp}/output_mixfolder', image_name)
            paths.append(file_path)
        return paths
    elif model == "nnunet":
        paths = []
        data = {"user_string": timestamp, "base_image_name": base_image_name}
        data1 = {"user_string": timestamp}
        response = requests.post("http://114.55.245.149:7833/mix_images/", params=data)
        #response = requests.get("http://114.55.245.149:7833/list_images/", params=data1)
        img_list = response.text
        image_list = ast.literal_eval(img_list)
        for image_name in image_list:
            file_path = os.path.join(f'http://114.55.245.149:7833/colorinput/{timestamp}/output_mixfolder', image_name)
            paths.append(file_path)
        return paths


with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.HTML("""
    <h1 style="text-align: center;">OCT层厚分析</h1>
    """)
    model = gr.Dropdown(
        ["DeeplabV3", "nnunet"], label="选择模型"
    )
    output2 = gr.Gallery(label="Layer Images", object_fit="scale-down", columns=4)
    with gr.Row():
        input = gr.Image(label="Input", sources=["upload"], height=200, width=200)
        output = gr.Gallery(label="calculate result", object_fit="scale-down", height=200)
        with gr.Column():
            btn1 = gr.Button("Upload")
            btn2 = gr.Button("Process Image")
            btn3 = gr.Button("Show Layer Images")
            btn4 = gr.ClearButton([input, output, output2], value="Reset")
        timestamp = gr.Textbox(visible=False)
        path = gr.Textbox(visible=False)
    btn1.click(fn=upload_image, inputs=[model, input], outputs=[timestamp, path])
    btn2.click(fn=process_img, inputs=[model, timestamp, path], outputs=output)
    btn3.click(fn=mix_img, inputs=[model, timestamp, path], outputs=output2)

demo.launch(server_name="0.0.0.0")

